from __future__ import annotations
from typing import Type, Dict, Optional, List

import math
from common.geoimage.scene_meta import SceneMeta

from catalog.models.base_vector_item import BaseVectorItem, extract_base_info
from common.geometry.bbox import BoundingBox
from common.geometry.geom import Geom
from common.geometry.srs import epsg_to_wkt, mgrs_tile_to_epsg
from peewee import CharField, FloatField

MGRS_GPKG_PATH = "/NAS/datasets/Places/mgrs_esa_100km_tile.gpkg"


class MGRSRecord(BaseVectorItem):
    collection_id = "ESA/MGRS"

    region_type = CharField(max_length=4)  # tile:48RVV or gzd:48R

    bbox_left_extended_utm = FloatField(index=True)
    bbox_right_extended_utm = FloatField(index=True)
    bbox_upper_extended_utm = FloatField(index=True)
    bbox_bottom_extended_utm = FloatField(index=True)

    @classmethod
    def create_one_item(cls: Type[MGRSRecord], info_item: Dict):
        base_info = extract_base_info(info_item)
        region_type = info_item["properties"]["region_type"]
        utm_wkt_esa = info_item["properties"]["geom_utm_esa"]
        geom_esa_utm_bounds = Geom.from_wkt(utm_wkt_esa).bounds
        bound_esa_utm_left = geom_esa_utm_bounds.left
        bound_esa_utm_right = geom_esa_utm_bounds.right
        bound_esa_utm_upper = geom_esa_utm_bounds.upper
        bound_esa_utm_bottom = geom_esa_utm_bounds.bottom

        return cls(
            region_type=region_type,
            bbox_left_extended_utm=bound_esa_utm_left,
            bbox_right_extended_utm=bound_esa_utm_right,
            bbox_upper_extended_utm=bound_esa_utm_upper,
            bbox_bottom_extended_utm=bound_esa_utm_bottom,
            **base_info,
        )

    def get_scenemeta(self, resolution=10, n_bands=1):
        meta_dict = self.get_mgrs_metadict_with_resolution(resolution)
        meta_dict["n_bands"] = n_bands
        scene_meta = SceneMeta.from_dict(meta_dict)
        return scene_meta

    def get_mgrs_metadict_with_resolution(self, resolution=10):
        supported_resolution = {10, 20, 60}
        if resolution not in supported_resolution:
            raise ValueError(
                "Resolution {} is not supported for mgrs meta calculation. Please select from {}".format(
                    resolution, supported_resolution
                )
            )
        # calculate rows & cols:
        n_rows = math.ceil(
            abs(self.bbox_upper_extended_utm - self.bbox_bottom_extended_utm)
            / resolution
        )

        n_cols = math.ceil(
            abs(self.bbox_right_extended_utm - self.bbox_left_extended_utm) / resolution
        )

        # extract prj_wkt
        epsg = mgrs_tile_to_epsg(self.uid)
        prj_wkt = epsg_to_wkt(epsg)

        # build geotrans
        geo_trans = (
            self.bbox_left_extended_utm,
            resolution,
            0,
            self.bbox_upper_extended_utm,
            0,
            -resolution,
        )

        meta_dict = {
            "n_cols": n_cols,
            "n_rows": n_rows,
            "n_bands": 1,
            "prj_wkt": prj_wkt,
            "geotrans": geo_trans,
            "value_interpretations": None,
        }

        return meta_dict

    @classmethod
    @BaseVectorItem.auto_connection_value
    def query_many_items(
        cls,
        spatial: Optional[Geom] = None,
        include_invalidated: bool = False,
    ) -> List[MGRSRecord]:
        query, _ = cls.construct_query(
            spatial=spatial, include_invalidated=include_invalidated
        )

        result_list = [item for item in query]
        return result_list


def ingest(mgrs_path=MGRS_GPKG_PATH):
    """
    This function is used to ingest geo information from MGRS_GPKG_PATH to the PostGreSQL

    Parameters
    ----------
    mgrs_path : str, optional
        path to the mgrs_gpkg file, by default MGRS_GPKG_PATH
    """
    import geopandas as gpd
    import shapely
    from tqdm import tqdm as tqdm

    margin_size = 4900

    gdf = gpd.read_file(mgrs_path)
    pbar = tqdm(total=len(gdf))
    for row in gdf.itertuples():
        pbar.update(1)
        geom_esa_utm_bounds = Geom.from_wkt(row.UTM_WKT).bounds
        bound_esa_utm_left = geom_esa_utm_bounds.left
        bound_esa_utm_right = geom_esa_utm_bounds.right
        bound_esa_utm_upper = geom_esa_utm_bounds.upper
        bound_esa_utm_bottom = geom_esa_utm_bounds.bottom

        bound_utm_left = bound_esa_utm_left + margin_size
        bound_utm_right = bound_esa_utm_right - margin_size
        bound_utm_upper = bound_esa_utm_upper - margin_size
        bound_utm_bottom = bound_esa_utm_bottom + margin_size

        geom_bounds = BoundingBox(
            bound_utm_left, bound_utm_bottom, bound_utm_right, bound_utm_upper
        )
        geom_utm = Geom.from_bbox(geom_bounds)
        geom_utm_wkt = geom_utm.to_wkt()

        tmp_info_item = {
            "properties": {
                "region_type": "tile",
                "geom_utm_esa": row.UTM_WKT,
                "geom_utm": geom_utm_wkt,
            },
            "id": row.name,
            "geometry": shapely.geometry.mapping(row.geometry),
        }
        table_item = MGRSRecord.create_one_item(tmp_info_item)
        MGRSRecord.insert_one_item(table_item, on_conflict="update")
    print("Data ingestion finished")


if __name__ == "__main__":
    import fire

    fire.Fire(
        {
            "make-table": MGRSRecord.make_table,
            "ingest": ingest,
            "remove-table": MGRSRecord.remove_table,
        }
    )
